using Mono.Cecil.Cil;

namespace Cil.Ast;

public class IlAstBranchesParser : IlAstParserBase
{
    private readonly Dictionary<Instruction, Instruction> _incoming = new Dictionary<Instruction, Instruction>();
    private readonly List<Instruction> _processed = new List<Instruction>();

    #region Step 0: Method root

    /// <summary>
    /// Light method to help to make base nodes for method root.
    /// </summary>
    public GroupExpression ParseMethodRoot(Context context)
    {
        var group = new GroupExpression(context);
        var node = new UnparsedExpression(context, context.Method.Body.Instructions.First(), context.Method.Body.Instructions.Last());
        group.Add(node);
        return group;
    }

    #endregion

    #region Step 1: Parsing all try-catch-finally blocks

    /// <summary>
    /// Parses all Tcf blocks. Should be called after ParseBranches.
    /// 解析所有Tcf块。应该在ParseBranches之后调用。
    /// </summary>
    /// <param name="context"></param>
    public void ParseTcfBlocks(Context context)
    {
        var cecilHandlers = context.Method.Body.ExceptionHandlers;

        // Groups of handlers - when one Try have many Catches or/and Finally block
        // 处理程序组——当一个Try有多个catch或/and Finally块时
        var groups = cecilHandlers.GroupBy(item => new { item.TryStart, item.TryEnd }).OrderByDescending(gr => gr.Key.TryEnd.Offset);
        var root = MethodBodyRoot;

        foreach (var grouping in groups)
        {
            var tryinfo = grouping.First();

            // split if needed, and return Expression, that covers TCF block
            var tcfblock = SplitGroupAndReplace(context, tryinfo.TryStart, grouping.Last().HandlerEnd.Previous, groupit: true);
            var tryblock = SplitGroupAndReplace(context, tryinfo.TryStart, tryinfo.TryEnd.Previous, groupit: true, exp_root: tcfblock);

            Expression faultBlock = null;
            Expression finallyBlock = null;
            Expression filterBlock = null;
            List<CatchBlockExpression> catches = new List<CatchBlockExpression>();

            foreach (var handler in grouping)
            {
                var group = SplitGroupAndReplace(context, handler.HandlerStart, handler.HandlerEnd.Previous,
                    exp_root: tcfblock);

                switch (handler.HandlerType)
                {
                    case ExceptionHandlerType.Catch:
                        var catchExpression = new CatchBlockExpression(context, handler.CatchType, group);
                        catches.Add(catchExpression);

                        // looking up first instructions list
                        var expression = group.FindFirstPostorder<UnparsedExpression>();

                        // looking up instruction, which reads exception from stack
                        var instruction = ResolveStackBlockEnd(expression.Instructions.First, -1);
                        if (instruction != null)
                        {
                            catchExpression.VariableReference = (instruction.Operand as VariableReference);
                            context.UserLocals[catchExpression] = instruction;
                        }

                        break;

                    case ExceptionHandlerType.Filter:
                        filterBlock = group;
                        break;

                    case ExceptionHandlerType.Finally:
                        finallyBlock = group;
                        break;

                    case ExceptionHandlerType.Fault:
                        faultBlock = group;
                        break;
                }
            }

            // create uncreated
            if (filterBlock == null) filterBlock = new Expression(context);
            if (finallyBlock == null) finallyBlock = new Expression(context);
            if (faultBlock == null) faultBlock = new Expression(context);

            // replace with our result
            var tcf = new TryExpression(context, grouping.ToList(), tryblock, faultBlock, filterBlock, finallyBlock, catches.ToArray());
            tcfblock.ReplaceWith(tcf);
        }
    }

    /// <summary>
    /// Should be used to split abstract instructions range.
    /// Returns Expression, that covers given instructions range
    /// 
    /// 应该用来分割抽象指令的范围。
    /// 返回表达式，它涵盖给定的指令范围
    /// </summary>
    private Expression SplitGroupAndReplace(Context context, Instruction start, Instruction end, bool groupit = false, Expression exp_root = null)
    {
        var gr_left = context.GetExpression(start);
        var gr_right = context.GetExpression(end);
        var root = exp_root ?? GetRootForNodes(gr_left, gr_right);

        if (root == null)
        {
            throw new ArgumentException("Given instructions have different roots");
        }

        gr_left = LookupNearestSubnode(root, gr_left);
        gr_right = LookupNearestSubnode(root, gr_right);

        // if both sides in one node
        // 如果两边都在一个节点中
        if (gr_left == gr_right)
        {
            // if this node isn't parsed 
            if (gr_left is UnparsedExpression)
            {
                // replace node with group of three subnodes, where
                // one subnode contains our instructions range.
                // 将node替换为一组三个子节点，其中一个子节点包含我们的指令范围。
                var node = gr_left as UnparsedExpression;
                Expression coverto;
                var last = node.Instructions.Last;
                var group = node.ParentNode as GroupExpression;
                var index = group.Nodes.IndexOf(node);

                if (start != node.Instructions.First)
                {
                    node.Instructions.Last = start.Previous;
                    index++;
                }
                else
                    group.RemoveAt(index);

                var item = new UnparsedExpression(context, start, end);
                if (groupit)
                    group.Insert(index, coverto = new GroupExpression(context) { item });
                else
                    group.Insert(index, coverto = item);
                index++;

                if (end != node.Instructions.Last)
                    group.Insert(index, new UnparsedExpression(context, end.Next, last));

                return coverto;
            }

            // node is parsed (may be if-else): nothing to split
            // Node被解析(可能是if-else):没有要拆分的东西
            return gr_left;
        }
        else
        {
            if (!(gr_left.ParentNode is GroupExpression))
                throw new ArgumentException("nodes root should be group");

            var rootgroup = gr_left.ParentNode as GroupExpression;
            var group = new GroupExpression(context);

            var i_left = rootgroup.IndexOf(gr_left);
            var i_right = rootgroup.IndexOf(gr_right);

            // parse left side if unparsed or include it whole
            // 如果未解析，则解析左侧，或完整包含左侧
            if (gr_left is UnparsedExpression && (gr_left as UnparsedExpression).Instructions.First != start)
            {
                var node = gr_left as UnparsedExpression;
                node.Instructions.Last = start.Previous;
                group.Add(new UnparsedExpression(context, start, node.Instructions.Last));
                node.ResetInstructionsInMap();
            }
            else
            {
                rootgroup.Remove(gr_left);
                group.Add(gr_left);
            }

            // move all between side expressions into group
            // 将所有边表达式之间移动到组中
            while (i_right != i_left + 1)
            {
                group.Add(rootgroup[i_left + 1]);
                rootgroup.RemoveAt(i_left + 1);
                i_right--;
            }

            if (gr_right is UnparsedExpression && (gr_right as UnparsedExpression).Instructions.Last != end)
            {
                var node = gr_right as UnparsedExpression;
                group.Add(new UnparsedExpression(context, node.Instructions.First, end));
                node.Instructions.First = end;
                node.ResetInstructionsInMap();
            }
            else
            {
                rootgroup.Remove(gr_right);
                group.Add(gr_right);
            }

            // insert all grouped to the root
            // 插入所有分组到根目录
            rootgroup.Insert(i_left, group);

            return group;
        }
    }

    /// <summary>
    /// Looks up by parents to root's child in the same subtree.
    /// 在同一子树中由父查找根的子。
    /// </summary>
    private Expression LookupNearestSubnode(Expression root, Expression node)
    {
        while (node.ParentNode != root)
            node = node.ParentNode;
        return node;
    }

    /// <summary>
    /// Looks for nearest to leafes parent node
    /// 查找离叶节点最近的父节点
    /// </summary>
    private Expression GetRootForNodes(Expression left, Expression right)
    {
        var parents = new List<Expression>();
        bool first = true;
        while (right != null)
        {
            if (!first) parents.Add(right);
            first = false;
            right = right.ParentNode;
        }

        first = true;
        while (left != null)
        {
            if (!first && parents.Contains(left))
            {
                return left;
            }

            first = false;
            left = left.ParentNode;
        }

        return null;
    }

    #endregion

    #region Step 2a: For if-else-while: calculate which instructions are the targets of jump instructions.

    public void ParseBranches(Context context, List<UnparsedExpression> items)
    {
        CalculateIncomingJumps(context.Method.Body);
        foreach (var item in items)
        {
            ParseInstructionsBlock(context, item);
        }
    }

    #endregion

    #region Step 2b: For if-else-while: Walking from last to first instruction, will find all branches

    private void ParseInstructionsBlock(Context context, UnparsedExpression expression)
    {
        var from = expression.Instructions.First;
        var to = expression.Instructions.Last;

        var current = from;
        while (current != to)
        {
            Instruction found = null;
            Instruction target = current.Operand as Instruction;
            var isWhileViaDoWhile = (HasOutgoingJump(current) && (current.OpCode.Code == Code.Br) &&
                                     FindFlowControl(target, to, from, target.Previous, out found, true, FlowControl.Cond_Branch));

            if (current.OpCode.Code == Code.Switch)
            {
                ParseSwitch(context, from, to, expression, current);
            }

            if ((HasIncomingJump(current) && !_processed.Contains(_incoming[current])) || (found != null)) // Has incoming branch. This means we have "while"
            {
                ParseLoop(context, from, to, expression, current, isWhileViaDoWhile);
                return;
            }

            if (HasOutgoingJump(current) && !_processed.Contains(current)) // Has outgoing branch. This means we have "if-else" statement
            {
                ParseIfElse(context, from, to, expression, current);
                return;
            }

            current = current.Next;
        }
    }

    private void ParseSwitch(Context context, Instruction from, Instruction to, Expression expression,
        Instruction current)
    {
        var cases = new List<CaseExpression>();
        var instructions = current.Operand as Instruction[];
        if (instructions == null || instructions.Length < 2)
        {
            throw new ArgumentException("strange array: < 2");
        }

        // find first non-equal branch
        var index = 0;
        while (instructions[index].Offset == instructions[0].Offset)
            index++;

        // lookup jump from case to the end of switch
        Instruction found = null;
        if (FindFlowControl(current, instructions[index].Previous, instructions[index], to, out found, false))
        {
            UnparsedExpression postfix = null;
            var condStart = ResolveStackBlockStart(current);

            // in case we're at end of if-else true block and switch end points out of if-else
            var switchEnd = found.Offset > to.Offset ? to : found.Operand as Instruction;

            // default branch
            Instruction defaultStart = null;
            Instruction defaultEnd;
            Expression defaultNode = null;
            if (FindFlowControl(current, instructions[0].Previous, instructions[0], to, out defaultEnd, true))
            {
                _processed.Add(defaultEnd);
                if (defaultEnd.Previous != current)
                {
                    defaultStart = current.Next;
                    defaultNode = new UnparsedExpression(context, defaultStart, defaultEnd);
                }
                else
                {
                    if (defaultEnd.Operand != switchEnd)
                    {
                        defaultStart = defaultEnd.Operand as Instruction;
                        defaultEnd = switchEnd.Previous;
                        defaultNode = new UnparsedExpression(context, defaultStart, defaultEnd);
                    }
                    else
                    {
                        defaultNode = new Expression(context);
                    }
                }
            }

            // lookup cases
            var groups = instructions.Where(ins => ins != defaultStart).Select((ins, ind) => new { Instruction = ins, Index = ind })
                .GroupBy(gr => gr.Instruction.Offset);

            foreach (var @group in groups)
            {
                var blockStart = @group.First().Instruction;

                if (blockStart != switchEnd)
                {
                    // lookup case block end
                    var blockEnd = blockStart;
                    while ((blockEnd.OpCode.Code != Code.Br && (blockEnd.Operand != switchEnd)) &&
                           blockEnd.Next != switchEnd)
                    {
                        blockEnd = blockEnd.Next;
                    }

                    _processed.Add(blockEnd);
                    var block = new CaseExpression(context, group.Select(g => g.Index).ToArray(),
                        new UnparsedExpression(context, blockStart, blockEnd));
                    cases.Add(block);
                }
            }

            var switchGroup = new GroupExpression(context);
            if (from != condStart)
            {
                switchGroup.Add(new UnparsedExpression(context, from, condStart.Previous));
            }

            var conditionNode = new UnparsedExpression(context, condStart, current);

            switchGroup.Add(new SwitchExpression(context, current, conditionNode, defaultNode, cases.ToArray()));

            if (to != switchEnd)
            {
                switchGroup.Add(postfix = new UnparsedExpression(context, switchEnd, to));
            }

            expression.ReplaceWith(switchGroup);

            foreach (var @case in cases)
            {
                ParseInstructionsBlock(context, @case.Body as UnparsedExpression);
            }

            if (postfix != null) ParseInstructionsBlock(context, postfix);
            if (defaultNode is UnparsedExpression) ParseInstructionsBlock(context, defaultNode as UnparsedExpression);
            return;
        }

        // else
        throw new ArgumentException("Bad switch block");
    }

    private void ParseLoop(Context context, Instruction from, Instruction to, UnparsedExpression expression, Instruction current, bool doWhileJump)
    {
        var group = new GroupExpression(context);
        UnparsedExpression prefix = null, conditionResult, body, postfix = null;
        ConditionExpression condition;
        LoopType looptype;

        // we have while(){ ... } which is made via do{ ... } while(); template
        if (doWhileJump)
        {
            // @current points to unconditional jump to condition block.
            // condition block is placed after body, like in do .. while.
            if (from.Offset < current.Offset)
                prefix = new UnparsedExpression(context, from, current.Previous);

            body = new UnparsedExpression(context, current.Next, (current.Operand as Instruction).Previous);

            // condition
            var conditionEnd = _incoming[body.Instructions.First];
            var conditionStart = ResolveStackBlockStart(conditionEnd);
            conditionResult = new UnparsedExpression(context, conditionStart, conditionEnd);

            var conditionEvalStart = body.Instructions.Last.Next;
            var conditionEvalEnd = conditionStart.Previous;

            if (conditionEvalEnd.Offset >= conditionEvalStart.Offset)
            {
                var conditionEvaluation = new UnparsedExpression(context, body.Instructions.Last.Next, conditionStart.Previous);
                condition = new ConditionExpression(context, conditionResult, conditionEvaluation);
            }
            else
            {
                condition = new ConditionExpression(context, conditionResult);
            }

            if (conditionEnd.Offset < to.Offset)
                postfix = new UnparsedExpression(context, conditionEnd.Next, to);

            _processed.Add(conditionEnd);
            _processed.Add(current);

            looptype = LoopType.While;
        }
        // if loop ends with unconditional branch, we have "while(<condition>) { <body> }"
        else
        {
            var incoming = _incoming[current];
            _processed.Add(incoming);
            if (incoming.OpCode.FlowControl == FlowControl.Branch)
            {
                Instruction conditionEnd;
                if (!FindFlowControl(current, incoming, incoming.Next, to, out conditionEnd, true, FlowControl.Cond_Branch))
                {
                    throw new ArgumentException("looks like loop with no exit");
                }

                var conditionStart = ResolveStackBlockStart(conditionEnd);

                if (from.Offset < conditionStart.Offset)
                    prefix = new UnparsedExpression(context, from, conditionStart.Previous);

                conditionResult = new UnparsedExpression(context, conditionStart, conditionEnd);

                var conditionEvalStart = current;
                var conditionEvalEnd = conditionStart.Previous;

                if (conditionEvalEnd.Offset > conditionEvalStart.Offset)
                {
                    condition = new ConditionExpression(context, conditionResult, new UnparsedExpression(context, conditionEvalStart, conditionEvalEnd));
                }
                else
                {
                    condition = new ConditionExpression(context, conditionResult);
                }

                body = new UnparsedExpression(context, conditionEnd.Next, incoming);

                if (incoming.Offset < to.Offset)
                    postfix = new UnparsedExpression(context, incoming.Next, to);

                looptype = LoopType.While;
            }
            // otherwice we have " do { <body> } while(<condition>); "
            else
            {
                var conditionStart = ResolveStackBlockStart(incoming);

                if (from.Offset < current.Offset)
                    prefix = new UnparsedExpression(context, from, current.Previous);

                conditionResult = new UnparsedExpression(context, conditionStart, incoming);
                condition = new ConditionExpression(context, conditionResult);
                body = new UnparsedExpression(context, current, conditionStart.Previous);

                if (incoming.Offset < to.Offset)
                    postfix = new UnparsedExpression(context, incoming.Next, to);

                looptype = LoopType.DoWhile;
            }
        }

        var condinstruction = conditionResult.Instructions.Last;
        if (prefix != null) group.Add(prefix);
        group.Add(new LoopExpression(context, doWhileJump ? current : null, condinstruction, looptype, condition, body));
        if (postfix != null) group.Add(postfix);
        expression.ReplaceWith(group);

        // run for each subblock
        ParseInstructionsBlock(context, body);
        if (postfix != null)
        {
            ParseInstructionsBlock(context, postfix);
        }

        return;
    }

    private void ParseIfElse(Context context, Instruction from, Instruction to, UnparsedExpression expression, Instruction current)
    {
        _processed.Add(current);
        var group = new GroupExpression(context);
        var target = (current.Operand as Instruction);
        var condStarts = ResolveStackBlockStart(current);

        // Prefix before condition
        if (condStarts != from)
        {
            group.Add(new UnparsedExpression(context, from, condStarts.Previous));
        }

        // condition block = condition except branching instruction
        var conditionNode = new UnparsedExpression(context, condStarts, current);
        var falseStart = current.Next;
        var falseEnd = target.Previous;
        var trueStart = target;
        var trueEnd = to;

        Instruction uncondJump;
        Expression trueNode, falseNode, postfixNode;

        if (FindFlowControl(falseStart, falseEnd, trueStart, trueEnd, out uncondJump, isforward: false))
        {
            trueEnd = (uncondJump.Operand as Instruction).Previous;
            falseNode = new UnparsedExpression(context, falseStart, falseEnd); // "false" branch
            trueNode = new UnparsedExpression(context, trueStart, trueEnd); // "true" branch
            postfixNode = new UnparsedExpression(context, trueEnd.Next, to); // out of "if-else" block
            _processed.Add(uncondJump);
        }
        else
        {
            falseNode = new UnparsedExpression(context, falseStart, falseEnd); // "false" branch
            trueNode = new Expression(context); // "true" branch
            postfixNode = new UnparsedExpression(context, trueStart, trueEnd); // out of "if-else" block
        }

        var ifelse =
            new IfElseExpression(context, current,
                new ConditionExpression(context, conditionNode),
                new GroupExpression(context) { trueNode },
                new GroupExpression(context) { falseNode }
            );
        group.Add(ifelse);
        group.Add(new GroupExpression(context) { postfixNode });
        expression.ReplaceWith(group);

        // run for each subblock

        var unparsedTrueNode = trueNode as UnparsedExpression;
        if (unparsedTrueNode != null) ParseInstructionsBlock(context, unparsedTrueNode);

        var unparsedFalseNode = falseNode as UnparsedExpression;
        if (unparsedFalseNode != null) ParseInstructionsBlock(context, unparsedFalseNode);

        ParseInstructionsBlock(context, postfixNode as UnparsedExpression);

        return;
    }

    #endregion

    #region tools

    private void CalculateIncomingJumps(MethodBody body)
    {
        _incoming.Clear();
        foreach (var instruction in body.Instructions.Where(HasOutgoingJump))
        {
            _incoming[(instruction.Operand as Instruction)] = instruction;
        }
    }

    private bool HasIncomingJump(Instruction position)
    {
        return _incoming.ContainsKey(position);
    }

    private bool HasOutgoingJump(Instruction position)
    {
        switch (position.OpCode.Code)
        {
            case Code.Br: // Unconditional
            case Code.Brfalse: // if false, 0, null
            case Code.Brtrue: // if true, <>0, <>null
            case Code.Beq: // if 2 values equal
            case Code.Bge: // if first >= second
            case Code.Bgt: // if first > second
            case Code.Ble: // if first <= second
            case Code.Blt: // if first < second
            case Code.Bne_Un: // if unsigned1 != unsigned2
            case Code.Bge_Un: // if unsigned >= unsigned2
            case Code.Bgt_Un: // if unsigned > unsigned2
            case Code.Ble_Un: // if unsigned <= unsigned2
            case Code.Blt_Un: // if unsigned < unsigned2
                return true;
            default:
                return false;
        }
    }

    private Instruction ResolveStackBlockStart(Instruction to)
    {
        Instruction current = to;
        var stack_delta = 0;
        do
        {
            stack_delta += current.StackDelta();
            current = current.Previous;
        } while (stack_delta != 0);

        return current.Next;
    }

    private Instruction ResolveStackBlockEnd(Instruction from, int stackSizeWeNeed)
    {
        var current = from;
        var stack_delta = 0;
        do
        {
            stack_delta += current.StackDelta(forward: true);
            if (stack_delta == stackSizeWeNeed)
                return current;

            current = current.Next;
        } while (current != null);

        return null;
    }

    private bool FindFlowControl(Instruction fromStart, Instruction fromEnd,
        Instruction toStart, Instruction toEnd,
        out Instruction found,
        bool isforward,
        FlowControl flowcontrol = FlowControl.Branch)
    {
        bool starts = false;
        var cur = isforward ? fromStart : fromEnd;
        var end = isforward ? fromEnd : fromStart;
        found = null;

        while (cur != end)
        {
            if (starts)
            {
                cur = isforward ? cur.Next : cur.Previous;
            }

            starts = true;
            if (cur.OpCode.FlowControl == flowcontrol)
            {
                var offset = (cur.Operand as Instruction).Offset;
                if (toStart.Offset <= offset && offset <= toEnd.Offset)
                {
                    found = cur;
                    return true;
                }
            }
        }

        return false;
    }

    #endregion
}